package Algorithm.dynamicProgramming.treeDp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * 834. 树中距离之和 https://leetcode.cn/problems/sum-of-distances-in-tree/
 * 题目简述：求树中每个节点到其它所有节点距离的总和
 */
public class SumOfDistancesInTree {

    /**
     * 思路：相比于父节点，子节点子树中所有节点距离减一，子树外所有节点距离增一。dfs统计根节点的距离总和，并在过程中通过后序统计每个节点子树的大小。
     * 最后再dfs递推每个节点的距离总和ans[child] = ans[root] - size[child] + (n - size[child]);
     * 关键就在于思考父节点答案到子节点答案的递推关系
     */
    int n;
    List<Integer>[] childsList;
    int[] ans, size;
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        this.n = n;
        childsList = new ArrayList[n];
        Arrays.setAll(childsList, i -> new ArrayList<Integer>());
        for (int[] edge : edges) {
            childsList[edge[0]].add(edge[1]);
            childsList[edge[1]].add(edge[0]);
        }
        ans = new int[n];
        size = new int[n];
        //dfs统计根节点的距离总和，并在过程中统计每个节点子树的大小
        dfs(0, -1, 0);
        //通过树上dp从根节点开始递推每个节点的距离总和
        dfs2(0, -1);
        return ans;
    }

    public void dfs(int root, int parent, int depth) {
        ans[0] += depth;
        size[root] = 1;
        for (int child : childsList[root]) {
            if (child != parent) {
                dfs(child, root, depth + 1);
                size[root] += size[child];//树大小即为1+所有子树大小，通过后序统计
            }
        }
    }

    /**
     * 相比于父节点，子节点子树中所有节点距离减一，子树外所有节点距离增一
     */
    public void dfs2(int root, int parent) {
        for (int child : childsList[root]) {
            if (child != parent) {
                ans[child] = ans[root] - size[child] + (n - size[child]);
                dfs2(child, root);
            }
        }
    }
}
